import pandas as pd
import os
from sklearn.model_selection import train_test_split


def load_toutiao_data(file_path):
    """
    加载并处理toutiao数据集，将原始标签映射为连续整数。

    参数:
        file_path (str): 数据集的TXT文件路径。

    返回:
        data (pd.DataFrame): 包含'label'和'text_a'两列的数据框，标签已转换为连续整数。
    """
    with open(file_path, 'r', encoding='utf-8') as file:
        lines = file.readlines()

    # 提取原始标签和文本内容，并生成标签映射字典
    raw_labels = []
    texts = []
    for line in lines:
        parts = line.strip().split('_!_')
        if len(parts) >= 4:
            raw_label = parts[1]  # 原始标签（如101、102）
            text = parts[3]
            raw_labels.append(raw_label)
            texts.append(text)

    # 生成标签映射：将原始标签（如101）映射为连续整数（如0,1,2...）
    unique_labels = sorted(list(set(raw_labels)))  # 确保顺序一致
    label_map = {label: idx for idx, label in enumerate(unique_labels)}

    # 转换标签并创建DataFrame
    labels = [label_map[label] for label in raw_labels]
    df = pd.DataFrame({
        'label': labels,
        'text_a': texts
    })

    return df, label_map  # 返回映射关系以便后续使用


def split_and_save(df, output_dir, test_size=0.1, dev_size=0.1, random_seed=42):
    """
    划分数据集并保存为TSV文件。

    参数:
        df (pd.DataFrame): 需要划分的数据框。
        output_dir (str): 输出目录路径。
        test_size (float): 测试集比例。
        dev_size (float): 验证集比例。
        random_seed (int): 随机种子。
    """
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"创建目录: {output_dir}")

    # 划分数据集
    train_dev_df, test_df = train_test_split(df, test_size=test_size,
                                             random_state=random_seed,
                                             stratify=df['label'])
    dev_size_relative = dev_size / (1 - test_size)
    train_df, dev_df = train_test_split(train_dev_df,
                                        test_size=dev_size_relative,
                                        random_state=random_seed,
                                        stratify=train_dev_df['label'])

    # 保存为TSV文件
    train_df.to_csv(f"{output_dir}/train.tsv", sep='\t', index=False, header=True)
    dev_df.to_csv(f"{output_dir}/dev.tsv", sep='\t', index=False, header=True)
    test_df.to_csv(f"{output_dir}/test.tsv", sep='\t', index=False, header=True)


# 使用示例
if __name__ == "__main__":
    input_file = 'E:/词向量/dataset/中文多分类/toutiao-text-classfication-dataset-master（多分类）/toutiao_cat_data.txt'  # 替换为您的数据文件路径
    output_dir = './toutiao_tsv'  # 替换为输出目录

    # 加载数据并获取标签映射
    df, label_map = load_toutiao_data(input_file)
    print("标签映射关系：", label_map)  # 输出原始标签与新标签的对应关系

    # 划分并保存数据集
    split_and_save(df, output_dir, test_size=0.1, dev_size=0.1, random_seed=42)